import os

from PyQt5.QtCore import Qt
from PyQt5.QtSql import QSqlDatabase, QSqlQuery
from PyQt5.QtWidgets import QTreeWidget, QTreeWidgetItem


def openDB():
    db = QSqlDatabase.addDatabase("QSQLITE")
    # db.setDatabaseName("../db/ProductManagement_new.db")
    db.setDatabaseName("./db/ProductManagement_new.db")  # 整合框架时使用
    db.open()
    return db


def delelePDF():
    # 删除指定目录中的PDF
    for parent, dirnames, filenames in os.walk('C:\\Users\\Public\\Documents\\PDF'):
        for fn in filenames:
            if fn.lower().endswith('.pdf'):
                os.remove(os.path.join(parent, fn))


def createModelTree():
    # 生成产品模型树
    modelTree = QTreeWidget()
    modelTree.header().hide()
    modelTree.setColumnCount(1)
    modelTree.setFocusPolicy(Qt.NoFocus)
    root = QTreeWidgetItem(modelTree)
    root.setText(0, "组件模型树")
    root.setText(1, "0")
    modelTree.addTopLevelItem(root)
    createModelTreeList(root)
    return modelTree


def createModelTreeList(root):
    # 生成产品模型树第一级列表
    result = getModelList()
    count = 0
    for i in result:
        childItem = QTreeWidgetItem(root)
        childItem.setText(0, i)
        childItem.setText(1, i)
        root.addChild(childItem)
        first_level = getTreeData(i)
        setTreeLevel(i, first_level, childItem, parent=i)
        count += 1


def getTreeData(modelName, fatherID="无父组件"):
    # 根据产品或组件得到子组件的数据
    db = openDB()
    query = QSqlQuery()
    sql = "SELECT ComponentName, ComponentID FROM T_Component WHERE ModelName= '%s' and fatherID= '%s'" % (
    modelName, fatherID)
    query.exec(sql)
    result = []
    while query.next():
        result.append([query.value(0), query.value(1)])
    return result


def setTreeLevel(modelName, current_level, root, parent='-'):
    for i in range(len(current_level)):
        levelOne = QTreeWidgetItem(root)
        levelOne.setText(0, str(current_level[i][0]))
        levelOne.setText(1, str(parent))

        second_level = getTreeData(modelName=modelName, fatherID=current_level[i][1])
        if second_level:
            setTreeLevel(modelName, second_level, levelOne, parent=parent)


def getModelList():
    # 得到所有的产品类型
    db = openDB()
    query = QSqlQuery()
    sql = "SELECT ProductName FROM T_Model"
    query.exec(sql)
    result = []
    while query.next():
        result.append(query.value(0))
    db.close()
    return result
